"""
SEGY到SEG2格式转换模块

该模块实现了将SEGY格式数据转换为SEG2格式的功能。
主要用于在生成SEGY文件的同时，自动生成同名的SEG2文件。

主要功能：
1. 读取SEGY数据文件
2. 提取元数据和波形数据
3. 将数据转换为SEG2格式
4. 保存SEG2文件

使用方法：
    from segy_to_seg2 import convert_segy_to_seg2
    
    # 转换单个SEGY文件
    seg2_file = convert_segy_to_seg2("path/to/segy_file.segy", "output_dir")
    
    # 批量转换目录中的SEGY文件
    seg2_files = convert_directory_segy_to_seg2("input_dir", "output_dir")
"""

import os
import sys
import struct
import numpy as np
from pathlib import Path
from datetime import datetime
from typing import List, Tuple, Dict, Optional, Union

# 使用ObsPy库处理SEGY格式
from obspy import read, Trace, Stream
from obspy.core import AttribDict
from obspy.core.utcdatetime import UTCDateTime

# 导入自定义SEG2写入器
from seg2_writer import write_seg2


def convert_segy_to_seg2(segy_file: Union[str, Path], 
                         output_dir: Optional[Union[str, Path]] = None,
                         output_file: Optional[Union[str, Path]] = None) -> Path:
    """
    将SEGY文件转换为SEG2格式
    
    参数:
        segy_file: SEGY文件路径
        output_dir: 输出目录（可选，如果不提供则使用SEGY文件所在目录）
        output_file: 输出文件路径（可选，如果不提供则使用与SEGY文件同名但扩展名为.seg2的文件）
        
    返回:
        生成的SEG2文件路径
    """
    segy_path = Path(segy_file)
    
    # 检查输入文件是否存在
    if not segy_path.exists():
        raise FileNotFoundError(f"SEGY文件不存在: {segy_path}")
    
    # 确定输出文件路径
    if output_file is not None:
        seg2_path = Path(output_file)
    else:
        if output_dir is not None:
            output_path = Path(output_dir)
            output_path.mkdir(parents=True, exist_ok=True)
        else:
            output_path = segy_path.parent
        
        # 使用与SEGY文件同名但扩展名为.seg2的文件名
        seg2_path = output_path / f"{segy_path.stem}.seg2"
    
    # 读取SEGY文件
    try:
        # 使用ObsPy读取SEGY文件
        st = read(str(segy_path), format="SEGY")
        
        # 检查是否成功读取
        if len(st) == 0:
            raise ValueError(f"SEGY文件不包含任何道: {segy_path}")
        
        # 使用自定义SEG2写入器保存为SEG2格式
        write_seg2(st, seg2_path)
        
        print(f"已成功将SEGY文件转换为SEG2格式: {seg2_path}")
        return seg2_path
        
    except Exception as e:
        raise RuntimeError(f"转换SEGY到SEG2失败: {e}")


def convert_directory_segy_to_seg2(input_dir: Union[str, Path],
                                  output_dir: Union[str, Path],
                                  recursive: bool = False) -> List[Path]:
    """
    批量转换目录中的SEGY文件为SEG2格式
    
    参数:
        input_dir: 输入目录
        output_dir: 输出目录
        recursive: 是否递归处理子目录
        
    返回:
        生成的SEG2文件路径列表
    """
    input_path = Path(input_dir)
    output_path = Path(output_dir)
    output_path.mkdir(parents=True, exist_ok=True)
    
    # 查找所有SEGY文件
    if recursive:
        segy_files = list(input_path.rglob("*.segy")) + list(input_path.rglob("*.SEGY"))
    else:
        segy_files = list(input_path.glob("*.segy")) + list(input_path.glob("*.SEGY"))
    
    if not segy_files:
        raise ValueError(f"在目录 {input_path} 中未找到SEGY文件")
    
    # 转换每个文件
    seg2_files = []
    for segy_file in segy_files:
        try:
            # 为每个文件创建相对路径
            relative_path = segy_file.relative_to(input_path)
            file_output_dir = output_path / relative_path.parent
            file_output_dir.mkdir(parents=True, exist_ok=True)
            
            # 转换文件
            seg2_file = convert_segy_to_seg2(segy_file, file_output_dir)
            seg2_files.append(seg2_file)
        except Exception as e:
            print(f"转换文件 {segy_file} 失败: {e}")
    
    return seg2_files


def convert_segy_to_seg2_same_dir(segy_file: Union[str, Path]) -> Path:
    """
    将SEGY文件转换为同目录下的SEG2文件
    
    参数:
        segy_file: SEGY文件路径
        
    返回:
        生成的SEG2文件路径
    """
    return convert_segy_to_seg2(segy_file)


def main():
    """命令行接口"""
    import argparse
    
    parser = argparse.ArgumentParser(description="SEGY到SEG2格式转换工具")
    parser.add_argument("input", help="输入SEGY文件或目录")
    parser.add_argument("--output", "-o", help="输出目录或文件")
    parser.add_argument("--recursive", "-r", action="store_true", help="递归处理子目录")
    
    args = parser.parse_args()
    
    input_path = Path(args.input)
    
    try:
        if input_path.is_file():
            # 单个文件转换
            if args.output:
                output_path = Path(args.output)
                if output_path.suffix:  # 如果有扩展名，视为文件
                    seg2_file = convert_segy_to_seg2(input_path, output_file=output_path)
                else:  # 否则视为目录
                    output_path.mkdir(parents=True, exist_ok=True)
                    seg2_file = convert_segy_to_seg2(input_path, output_path)
            else:
                # 输出到同一目录
                seg2_file = convert_segy_to_seg2_same_dir(input_path)
            
            print(f"转换完成: {seg2_file}")
            
        elif input_path.is_dir():
            # 目录批量转换
            if not args.output:
                # 如果没有指定输出目录，使用输入目录
                output_dir = input_path
            else:
                output_dir = Path(args.output)
            
            seg2_files = convert_directory_segy_to_seg2(
                input_path, output_dir, recursive=args.recursive
            )
            
            print(f"批量转换完成，共生成 {len(seg2_files)} 个SEG2文件:")
            for f in seg2_files:
                print(f"  {f}")
        
        else:
            print(f"错误: 输入路径不存在: {input_path}")
            return 1
        
        return 0
        
    except Exception as e:
        print(f"错误: {e}")
        return 1


if __name__ == "__main__":
    sys.exit(main())